#include <cstdio>
#include <systemc>
#include <tlm.h>
#include <tlm_utils/peq_with_cb_and_phase.h>
#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>

using namespace sc_core;
using namespace sc_dt;
using namespace std;

// User-defined Memory Manager, which maintains a pool of transactions
class MemoryManager : public tlm::tlm_mm_interface {
    typedef tlm::tlm_generic_payload gp_t;
public:
    MemoryManager() {}
    ~MemoryManager() {
        for (std::size_t i = 0; i < all_trans.size(); ++i) {
            delete all_trans[i];
        }
    }
    // Allocate a new transaction or reuse one from the free pool
    gp_t* allocate() {
        gp_t* ptr;
        if (!free_pool.empty()) {
            ptr = free_pool.back();
            free_pool.pop_back();
        } else {
            ptr = new gp_t(this);
            all_trans.push_back(ptr);
        }
        printf("------ Called MemoryManager.allocate(), trans_count = %zu\n", ++count);
        return ptr;
    }
    // Free a transaction back to the memory manager
    void free(gp_t* trans) {
        trans->reset();
        free_pool.push_back(trans);
        printf("------ Called MemoryManager.free(), trans_count = %zu\n", --count);
    }
private:
    std::vector<gp_t*> free_pool;
    std::vector<gp_t*> all_trans;
    size_t count = 0;
};

// Generate a random delay
int rand_ps() {
    int n = rand() % 100;
    n = n * n * n;
    return n / 100;
}


// **************************************************************************************
// Initiator module generating multiple pipelined generic payload transactions
// **************************************************************************************
SC_MODULE(Initiator) {
    // TLM-2.0 non-blocking initiator socket
    tlm_utils::simple_initiator_socket<Initiator> socket;

    SC_CTOR(Initiator)
        : socket("socket"), request_in_progress_ptr(nullptr), peq(this, &Initiator::peq_cb) {
        // Register the backward path callback function
        socket.register_nb_transport_bw(this, &Initiator::nb_transport_bw);

        SC_THREAD(thread_process);
    }

    void thread_process() {
        tlm::tlm_generic_payload *trans;
        tlm::tlm_phase phase;
        sc_time delay;

        // Generate a sequence of random transactions
        for (int i = 0; i < 1024; i++) {
            // Prepare the transaction attributes
            int addr = rand();
            tlm::tlm_command cmd = static_cast<tlm::tlm_command>(rand() % 2);
            if (cmd == tlm::TLM_WRITE_COMMAND) {
                data[i % 16] = rand();
            }

            // Grab a new transaction from the memory manager
            trans = mm.allocate();
            trans->acquire();

            // Set all attributes except byte_enable_length and extensions (unused)
            trans->set_command(cmd);
            trans->set_address(addr);
            trans->set_data_ptr(reinterpret_cast<unsigned char *>(&data[i % 16]));
            trans->set_data_length(4);
            trans->set_streaming_width(4); // Set to data_length to indicate no streaming
            trans->set_byte_enable_ptr(0); // 0 indicates unused
            trans->set_dmi_allowed(false); // Mandatory initial value
            trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE); // Mandatory initial value

            // The `BEGIN_REQ/END_REQ exclusion rule` ensures pipelined flow control.
            // A new BEGIN_REQ must not be sent until the END_REQ of the previous one is received.
            if (request_in_progress_ptr) {
                // If a request is already in flight, wait for its completion event.
                // This event is notified by peq_cb upon receiving an END_REQ from the target.
                wait(end_request_event);
            }
            // Set the flag to the current transaction, indicating a new request is now in flight.
            request_in_progress_ptr = trans;
            // Set the initial phase for the forward path call.
            phase = tlm::BEGIN_REQ;

            // Timing annotation: models the initiator's internal processing delay before the call.
            delay = sc_time(rand_ps(), SC_PS);

            printf("0x%x new, cmd=%c, data=%x at %s\n", addr, (cmd ? 'W' : 'R'), data[i % 16], sc_time_stamp().to_string().c_str());

            // Initiate the non-blocking forward path call.
            // It returns immediately with a synchronous status.
            tlm::tlm_sync_enum status = socket->nb_transport_fw(*trans, phase, delay);

            // In this specific example, the target always returns TLM_ACCEPTED, so these branches will not be taken.
            // Check value returned from nb_transport_fw
            if (status == tlm::TLM_UPDATED) {
                // The timing annotation must be honored
                peq.notify(*trans, phase, delay);
            }
            else if (status == tlm::TLM_COMPLETED) {
                // The completion of the transaction necessarily ends the BEGIN_REQ phase
                request_in_progress_ptr = nullptr;
                // The target has terminated the transaction
                check_transaction(*trans);
            }
            // Simulate some delay between transaction generations
            wait(sc_time(rand_ps(), SC_PS));
        }

        wait(100, SC_NS);

        // Allocate a transaction for one final, nominal call to b_transport
        trans = mm.allocate();
        trans->acquire();
        trans->set_command(tlm::TLM_WRITE_COMMAND);
        trans->set_address(0);
        trans->set_data_ptr(reinterpret_cast<unsigned char *>(&data[0]));
        trans->set_data_length(4);
        trans->set_streaming_width(4); // = data_length to indicate no streaming
        trans->set_byte_enable_ptr(0); // 0 indicates unused
        trans->set_dmi_allowed(false); // Mandatory initial value
        trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE); // Mandatory initial value

        delay = sc_time(rand_ps(), SC_PS);

        printf("Calling b_transport at %s with delay = %s\n", sc_time_stamp().to_string().c_str(), delay.to_string().c_str());

        // Call b_transport to demonstrate the b/nb conversion by the simple_target_socket
        socket->b_transport(*trans, delay);
        check_transaction(*trans);
    }

    // TLM-2.0 backward-path non-blocking transport method. Called by the Target.
    virtual tlm::tlm_sync_enum nb_transport_bw(tlm::tlm_generic_payload &trans,
                                               tlm::tlm_phase &phase, sc_time &delay) {
        // Received a phase/time update from the target.
        // Add it to the event queue for execution at the correct simulation time.
        peq.notify(trans, phase, delay);
        return tlm::TLM_ACCEPTED;
    }

    // The callback for the Payload Event Queue (PEQ), processing events from the Target at their scheduled time.
    void peq_cb(tlm::tlm_generic_payload &trans, const tlm::tlm_phase &phase) {
        if (phase == tlm::END_REQ) {
            printf("0x%llx END_REQ at %s\n", trans.get_address(), sc_time_stamp().to_string().c_str());
        }
        else if (phase == tlm::BEGIN_RESP) {
            printf("0x%llx BEGIN_RESP at %s\n", trans.get_address(), sc_time_stamp().to_string().c_str());
        }

        // The request phase has ended.
        // The end of the request phase is the precondition for the Initiator to be able to send a new request.
        if (phase == tlm::END_REQ ||
            // In this case, BEGIN_RESP will not start before END_REQ
            (&trans == request_in_progress_ptr && phase == tlm::BEGIN_RESP)) {
            // Clear flag, indicating the Initiator is ready to send a new request.
            request_in_progress_ptr = nullptr;
            // Notify the event to unblock the thread_process thread, which may be suspended in a wait()
            end_request_event.notify();
        }
        // An Initiator should never receive BEGIN_REQ or END_RESP on the backward path.
        else if (phase == tlm::BEGIN_REQ || phase == tlm::END_RESP) {
            SC_REPORT_FATAL("TLM-2", "Illegal transaction phase received by initiator");
        }

        // Received BEGIN_RESP from the Target
        if (phase == tlm::BEGIN_RESP) {
            // Check the transaction's response status
            check_transaction(trans);
            // TLM-2.0 PROTOCOL: After receiving BEGIN_RESP,
            // the Initiator must send END_RESP to complete the four-phase handshake.
            tlm::tlm_phase fw_phase = tlm::END_RESP;
            sc_time delay = sc_time(rand_ps(), SC_PS);
            socket->nb_transport_fw(trans, fw_phase, delay);
            // Ignore return value
        }
    }

    // Called on receiving BEGIN_RESP or TLM_COMPLETED
    void check_transaction(tlm::tlm_generic_payload &trans) {
        if (trans.is_response_error()) {
            char msg[128];
            sprintf(msg, "Transaction returned with error, response status = %s", trans.get_response_string().c_str());
            SC_REPORT_ERROR("TLM-2", msg);
        }

        tlm::tlm_command cmd = trans.get_command();
        sc_dt::uint64 addr = trans.get_address();
        int *ptr = reinterpret_cast<int *>(trans.get_data_ptr());

        printf("0x%x check, cmd=%c, data=%x at %s\n", (unsigned int)addr, (cmd ? 'W' : 'R'), *ptr, sc_time_stamp().to_string().c_str());

        // Allow the memory manager to free the transaction object
        trans.release();
    }

    MemoryManager mm;
    int data[16];
    tlm::tlm_generic_payload *request_in_progress_ptr;
    sc_event end_request_event;
    tlm_utils::peq_with_cb_and_phase<Initiator> peq;
};


// **************************************************************************************
// Target module able to handle two pipelined transactions
// **************************************************************************************
DECLARE_EXTENDED_PHASE(RWOP); // A user-defined phase for read/write operation
SC_MODULE(Target) {
    tlm_utils::simple_target_socket<Target> socket;

    SC_CTOR(Target)
        : socket("socket"), n_trans(0), response_in_progress(false),
          next_response_pending_ptr(nullptr), end_req_pending_ptr(nullptr),
          peq(this, &Target::peq_cb) {
        // Register callbacks for incoming interface method calls
        socket.register_nb_transport_fw(this, &Target::nb_transport_fw);
    }

    // TLM-2.0 forward-path non-blocking transport method. Called by the Initiator.
    virtual tlm::tlm_sync_enum nb_transport_fw(tlm::tlm_generic_payload &trans,
                                               tlm::tlm_phase &phase, sc_time &delay) {
        sc_dt::uint64 addr = trans.get_address();
        unsigned int len = trans.get_data_length();
        unsigned char *byte = trans.get_byte_enable_ptr();
        unsigned int width = trans.get_streaming_width();

        // Obliged to check the transaction attributes for unsupported features
        // and to generate the appropriate error response
        if (byte != 0) {
            trans.set_response_status(tlm::TLM_BYTE_ENABLE_ERROR_RESPONSE);
            return tlm::TLM_COMPLETED;
        }
        if (len > 4 || width < len) {
            trans.set_response_status(tlm::TLM_BURST_ERROR_RESPONSE);
            return tlm::TLM_COMPLETED;
        }

        // Put the transaction into the PEQ with the delay
        peq.notify(trans, phase, delay);

        // Immediately return TLM_ACCEPTED, indicating the request is accepted and future communication will be on the backward path.
        return tlm::TLM_ACCEPTED;
    }

    // The callback for the Payload Event Queue (PEQ), processing events from the Initiator at their scheduled time.
    void peq_cb(tlm::tlm_generic_payload &trans, const tlm::tlm_phase &phase) {
        tlm::tlm_sync_enum status;
        sc_time delay;

        switch (phase) {

        case tlm::BEGIN_REQ:
            printf("0x%llx BEGIN_REQ at %s\n", trans.get_address(), sc_time_stamp().to_string().c_str());
            // Increment the transaction reference count
            trans.acquire();
            // If the pipeline is full, apply back-pressure:
            //     1. Store the current request in a pending pointer
            //     2. Withhold the END_REQ response to prevent the Initiator from sending more requests
            if (n_trans == 2) {
                end_req_pending_ptr = &trans;
            }
            // If the pipeline has capacity
            else {
                // process the request immediately
                status = send_end_req(trans);
                // The transaction was aborted by the Initiator
                if (status == tlm::TLM_COMPLETED) break;
            }
            break;

        case tlm::END_RESP:
            // On receiving END_RESP, the Target can release the transaction and allow other pending transactions to proceed
            printf("0x%llx END_RESP at %s\n", trans.get_address(), sc_time_stamp().to_string().c_str());
            // TLM-2.0 PROTOCOL: Receiving END_RESP without a preceding BEGIN_RESP is a protocol violation.
            if (!response_in_progress) {
                SC_REPORT_FATAL("TLM-2", "Illegal transaction phase END_RESP received by target");
            }
            // The transaction is complete, so the Target releases its ownership.
            trans.release();
            n_trans--;
            // The response path is now free for the next response.
            response_in_progress = false;
            // Check for and process previously blocked responses and requests.
            if (next_response_pending_ptr) {
                send_response(*next_response_pending_ptr);
                next_response_pending_ptr = nullptr;
            }
            if (end_req_pending_ptr) {
                status = send_end_req(*end_req_pending_ptr);
                end_req_pending_ptr = nullptr;
            }
            break;

        case tlm::END_REQ:
        case tlm::BEGIN_RESP:
            SC_REPORT_FATAL("TLM-2", "Illegal transaction phase received by target");
            break;
        default:
            if (phase != RWOP) break;
            // Execute the read or write command
            tlm::tlm_command cmd = trans.get_command();
            sc_dt::uint64 addr = trans.get_address();
            unsigned char *ptr = trans.get_data_ptr();
            unsigned int len = trans.get_data_length();
            if (cmd == tlm::TLM_READ_COMMAND) {
                *reinterpret_cast<int *>(ptr) = rand(); // Random fill for read
                printf("0x%llx Execute READ, data = %x\n", addr, *reinterpret_cast<int *>(ptr));
            } else if (cmd == tlm::TLM_WRITE_COMMAND) {
                printf("0x%llx Execute WRITE, data = %x\n", addr, *reinterpret_cast<int *>(ptr));
            }
            // Set the response status to indicate successful completion
            trans.set_response_status(tlm::TLM_OK_RESPONSE);

            // Target must honor BEGIN_RESP/END_RESP exclusion rule
            // i.e. must not send BEGIN_RESP until receiving previous END_RESP or BEGIN_REQ
            if (response_in_progress) {
                // Target allows only two transactions in-flight
                // Response path is busy, so store this completed transaction and wait for the path to be free.
                if (next_response_pending_ptr) {
                    SC_REPORT_FATAL("TLM-2", "Attempt to have two pending responses in target");
                }
                next_response_pending_ptr = &trans;
            }
            // Response path is free, send the response now
            else {
                send_response(trans);
            }
            break;
        }
    }

    tlm::tlm_sync_enum send_end_req(tlm::tlm_generic_payload &trans) {
        tlm::tlm_sync_enum status;
        tlm::tlm_phase bw_phase;
        tlm::tlm_phase rwop_phase = RWOP;
        sc_time delay;

        // Send END_REQ to the Initiator, indicating the end of the request phase
        bw_phase = tlm::END_REQ;
        delay = sc_time(rand_ps(), SC_PS); // Accept delay
        status = socket->nb_transport_bw(trans, bw_phase, delay);
        if (status == tlm::TLM_COMPLETED) {
            // Transaction aborted by the initiator
            // TLM_UPDATED cannot occur at this point in the base protocol, so need not be checked
            trans.release();
            return status;
        }

        // Push the transaction into PEQ for the read/write operation phase
        delay = delay + sc_time(rand_ps(), SC_PS); // Latency
        peq.notify(trans, rwop_phase, delay);

        // Increment the count of transactions in the pipeline
        n_trans++;

        return status;
    }

    void send_response(tlm::tlm_generic_payload &trans) {
        tlm::tlm_sync_enum status;
        tlm::tlm_phase bw_phase;
        sc_time delay;

        response_in_progress = true; // Mark the response path as busy
        bw_phase = tlm::BEGIN_RESP;
        delay = SC_ZERO_TIME; // Immediate response
        status = socket->nb_transport_bw(trans, bw_phase, delay);

        // In this example, the initiator always returns TLM_ACCEPTED, so these branches will not be taken.
        if (status == tlm::TLM_UPDATED) {
            // The timing annotation must be honored
            peq.notify(trans, bw_phase, delay);
        } else if (status == tlm::TLM_COMPLETED) {
            // The initiator has terminated the transaction
            trans.release();
            n_trans--;
            response_in_progress = false;
        }
    }

    int n_trans;
    bool response_in_progress;
    tlm::tlm_generic_payload *next_response_pending_ptr;
    tlm::tlm_generic_payload *end_req_pending_ptr;
    tlm_utils::peq_with_cb_and_phase<Target> peq;
};

SC_MODULE(Top) {
    Initiator *initiator;
    Target *target;

    SC_CTOR(Top) {
        // Instantiate components
        initiator = new Initiator("initiator");
        target = new Target("target");

        // One initiator is bound directly to one target with no intervening bus
        // Bind initiator socket to target socket
        initiator->socket.bind(target->socket);
    }
};

int sc_main(int argc, char *argv[]) {
    Top top("top");
    sc_start();
    printf("Finished at %s\n", sc_time_stamp().to_string().c_str());
    return 0;
}
